# coding: utf-8

import matplotlib.pyplot as plt
import matplotlib
font = {"family": "Arial", "weight": "normal", "size": 14}
matplotlib.rc("font", **font)
from PIL import Image
import numpy as np
import glob
import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
from metrics import *


def plot_snr_test(origin_image_path, snr_image_folder, start, end, step=1, qam_order=16, save_path="D:/test.png"):
    ssims = []
    psnrs = []
    msssims = []
    snrs = list(range(start, end, step))

    for snr in snrs:
        origin_image = Image.open(origin_image_path)
        snr_image_name = os.path.basename(origin_image_path.replace(".png", "_%dqam_snr%d.png" % (qam_order, snr)))
        snr_image = Image.open(os.path.join(snr_image_folder, snr_image_name))
        ssim_val = compute_ssim(origin_image, snr_image)
        psnr_val = compute_psnr(origin_image, snr_image)
        msssim_val = compute_msssim(origin_image, snr_image)
        ssims.append(ssim_val)
        psnrs.append(psnr_val)
        msssims.append(msssim_val)

    print("SNR: ", snrs)
    print("PSNR: ", psnrs)
    print("SSIM: ", ssims)
    print("MS-SSIM: ", msssims)

    plt.figure(figsize=(18, 6))
    # plt.title("1/2LDPC+%dQAM+AWGN" % qam_order)
    plt.subplot(1, 3, 1)
    plt.plot(snrs, psnrs, "-", color="red", marker="^")
    plt.xlabel("SNR(dB)")
    plt.ylabel("PSNR(dB)")
    plt.xlim([start, end])
    plt.xticks(list(range(start, end, 2)))
    plt.grid()

    plt.subplot(1, 3, 2)
    plt.plot(snrs, ssims, "-", color="red", marker="^")
    plt.xlabel("SNR(dB)")
    plt.ylabel("SSIM")
    plt.xlim([start, end])
    plt.xticks(list(range(start, end, 2)))
    plt.grid()

    plt.subplot(1, 3, 3)
    plt.plot(snrs, msssims, "-", color="red", marker="^")
    plt.xlabel("SNR(dB)")
    plt.ylabel("MS-SSIM")
    plt.xlim([start, end])
    plt.xticks(list(range(start, end, 2)))
    plt.grid()
    # plt.show()
    plt.savefig(save_path, format="png", dpi=600)


if __name__ == '__main__':
    size = "512x768"
    origin_image_path = r"D:\video-communication-dataset\图像编解码测试\kodak\archive\kodim23.png"
    snr_image_folder = r"F:\PCL\month\2025\25-08\experiments\H.265\512x768\0.5LDPC16QAM"
    start = 5
    end = 26
    step = 2
    plot_snr_test(origin_image_path, snr_image_folder, start, end,
                  step=step, qam_order=16, save_path="D:/test.png")

